from osgeo import gdal
gdal.UseExceptions()
import numpy as np
import os
import math
from tqdm import tqdm

class CalcReliefAmplitude():
    def __init__(self, numrange, gridLength):
        self.numrange = numrange
        self.len = len(self.numrange)
        self.gridLength = gridLength

    def _calcVar(self):
        return np.var(self.numrange) * self.len

    def calcBestWindow(self):
        meanRAList = [math.log(a / math.pow((self.gridLength * (i + 2)), 2)) for i, a in enumerate(self.numrange)]
        S = []
        for j in range(2, len(meanRAList) + 1):
            s1 = np.var(meanRAList[:j-1]) * (j-1)
            s2 = np.var(meanRAList[j-1:]) * (len(meanRAList) - j + 1)
            S.append(s1 + s2)

        bestWindow = S.index(min(S)) + 1 + 2
        return [bestWindow, math.pow((bestWindow * self.gridLength), 2)]

def BlockStatistics(raster, stat_type):
    band = raster.GetRasterBand(1)
    array = band.ReadAsArray()
    nodata_value = band.GetNoDataValue()

    if nodata_value is not None:
        array = np.where(array == nodata_value, np.nan, array)  # 将NoData值设为NaN

    if stat_type == 'MAXIMUM':
        return np.nanmax(array)
    elif stat_type == 'MINIMUM':
        return np.nanmin(array)

def get_qfdRast(raster, i):
    band = raster.GetRasterBand(1)
    array = band.ReadAsArray()
    nodata_value = band.GetNoDataValue()

    if nodata_value is not None:
        array = np.where(array == nodata_value, np.nan, array)  # 将NoData值设为NaN

    # 在窗口大小范围内计算最大值和最小值
    window_size = i
    rows, cols = array.shape
    qfd_array = np.zeros((rows - window_size + 1, cols - window_size + 1))

    for r in range(rows - window_size + 1):
        for c in range(cols - window_size + 1):
            window = array[r:r + window_size, c:c + window_size]
            if np.all(np.isnan(window)):
                qfd_array[r, c] = np.nan  # 如果窗口内全是NaN值，设置为NaN
            else:
                qfd_array[r, c] = np.nanmax(window) - np.nanmin(window)

    return qfd_array

def getRasterList(raster_path):
    numrange = []
    rasterList = []
    raster = gdal.Open(raster_path)
    if raster is None:
        print("无法打开文件")
        return numrange, rasterList

    print(raster.GetProjection())
    print(raster.GetGeoTransform())

    for i in tqdm(range(2, 50), desc="Processing windows"):
        qfd = get_qfdRast(raster, i)
        qfdValue = np.nanmean(qfd) if qfdparams == 'mean' else np.nanmax(qfd)
        numrange.append(qfdValue)
        rasterList.append('dem_%s,%s' % (i, qfdValue))
        print('dem_%s,%s' % (i, qfdValue))
    return numrange, rasterList

if __name__ == '__main__':
    block = r'G:\2024工作\Dem\六石街道dem1_ProjectRaster1.tif'
    outputPath = r'G:\2024工作\dem_1'
    DEM_Length = 2  # DEM的精度，即格网大小，单位m
    qfdparams = 'mean'
    
    # 检查输出目录是否存在
    if not os.path.exists(outputPath):
        os.makedirs(outputPath)

    numrange, rasterList = getRasterList(block)
    bestWindow, bestArea = CalcReliefAmplitude(numrange, DEM_Length).calcBestWindow()
    output = os.path.join(outputPath, 'dem_%s' % bestWindow)
    with open(os.path.join(outputPath, 'qfd.csv'), 'a') as f:
        f.write('Raster,%s\n' % qfdparams)
        for line in rasterList:
            f.write(line + '\n')

    qfdResult = get_qfdRast(gdal.Open(block), bestWindow)
    qfdArray = np.array(qfdResult) if isinstance(qfdResult, (list, np.ndarray)) else np.array([qfdResult])

    # 调试信息
    print("qfdResult:", qfdResult)
    print("qfdArray:", qfdArray)
    print("qfdArray shape:", qfdArray.shape)

    if qfdArray.ndim == 2:
        driver = gdal.GetDriverByName("GTiff")
        outdata = driver.Create(output + '.tif', qfdArray.shape[1], qfdArray.shape[0], 1, gdal.GDT_Float32)
        outdata.GetRasterBand(1).WriteArray(qfdArray)
        outdata.FlushCache()
        print(u'The Best Window is %s*%s, Area is %s m²' % (bestWindow, bestWindow, bestArea))
    else:
        print("Error: qfdArray is not a 2D array.")
